import os, os.path as osp
import PIL.Image as Image
import cv2
import argparse
import numpy as np
import shutil
import subprocess

def cat_images(frame_dir, cam_img_size, pred_img_size, pred_img_size2, spacing):
    cam_order = [
        'CAM_FRONT_LEFT', 'CAM_FRONT', 'CAM_FRONT_RIGHT',
        'CAM_BACK_LEFT',  'CAM_BACK',  'CAM_BACK_RIGHT'
    ]
    pred_imgs = []
    orig_imgs = []
    for cam in cam_order:
        pred_path = osp.join(frame_dir, f'pred_{cam}.png')
        orig_path = osp.join(frame_dir, f'{cam}.jpg')
        if osp.exists(pred_path) and osp.exists(orig_path):
            pred_imgs.append(Image.open(pred_path).resize(cam_img_size, Image.BILINEAR))
            orig_imgs.append(Image.open(orig_path).resize(cam_img_size, Image.BILINEAR))
        else:
            print(f'警告: {pred_path} 或 {orig_path} 不存在，跳过该帧。')
            return None
    fn = osp.join(frame_dir, 'pred_normal.png')
    if osp.exists(fn):
        pred_normal = Image.open(fn).resize(pred_img_size, Image.BILINEAR)
    else:
        print(f"警告: {fn} 不存在，跳过该帧。")
        return None
    fn = osp.join(frame_dir, 'pred_bev.png')
    if osp.exists(fn):
        pred_bev = Image.open(fn).resize(pred_img_size, Image.BILINEAR).crop([460, 0, 1460, 1080])
    else:
        print(f"警告: {fn} 不存在，跳过该帧。")
        return None

    cam_w, cam_h = cam_img_size
    pred_w, pred_h = pred_img_size
    # 六列，两行（左3列原图，右3列pred）
    result_w = cam_w * 6 + 5 * spacing
    result_h = cam_h * 2 + pred_h + 2 * spacing
    result = Image.new(pred_imgs[0].mode, (result_w, result_h), (0, 0, 0))

    # 左侧三列 原图
    for i in range(3):
        result.paste(orig_imgs[i], box=((cam_w+spacing)*i, 0))
    for i in range(3, 6):
        result.paste(orig_imgs[i], box=((cam_w+spacing)*(i-3), cam_h+spacing))

    # 右侧三列 pred
    for i in range(3):
        result.paste(pred_imgs[i], box=((cam_w+spacing)*(i+3), 0))
    for i in range(3, 6):
        result.paste(pred_imgs[i], box=((cam_w+spacing)*((i-3)+3), cam_h+spacing))

    # 下方大图
    result.paste(pred_normal, box=(0, cam_h*2+spacing*2))
    result.paste(pred_bev, box=(pred_w+spacing, cam_h*2+spacing*2))

    result_path = osp.join(frame_dir, 'cat.png')
    result.save(result_path)

    return result

# 保证宽高为偶数
def pad_to_even(img):
    w, h = img.size
    new_w = w if w % 2 == 0 else w + 1
    new_h = h if h % 2 == 0 else h + 1
    if new_w != w or new_h != h:
        new_img = Image.new(img.mode, (new_w, new_h), (0, 0, 0))
        new_img.paste(img, (0, 0))
        return new_img
    return img

if __name__ == "__main__":
    parse = argparse.ArgumentParser('')
    parse.add_argument('--scene-dir', type=str, default='vis_results/scene/scene-630', 
                       help='directory of the scene outputs')
    args = parse.parse_args()

    scene_dir = args.scene_dir
    frame_dirs = [d for d in os.listdir(scene_dir) if osp.isdir(osp.join(scene_dir, d))]
    frame_dirs.sort(key=lambda x: int(x))
    cam_img_size = [480, 270]
    pred_img_size = [1920, 1080]
    pred_img_size2 = [1000, 1080]
    spacing = 10

    cat_imgs = []
    for dir in frame_dirs:
        frame_path = osp.join(scene_dir, dir)
        print(f'processing {frame_path}')
        cat_img = cat_images(frame_path, cam_img_size, pred_img_size, pred_img_size2, spacing)
        if cat_img is not None:
            cat_imgs.append(cat_img.resize((cat_img.width//2, cat_img.height//2), resample=Image.LANCZOS))

    if len(cat_imgs) == 0:
        print("没有可用的帧，未生成 GIF/视频。请检查数据是否完整。")
    else:
        cat_imgs[0].save(osp.join(scene_dir, 'video_resize.gif'),
                         save_all=True,
                         append_images=cat_imgs[1:],
                         optimize=True,
                         duration=150,
                         loop=0)
        # 保存视频（OpenCV低码率版本）
        video_path = osp.join(scene_dir, 'video_resize.mp4')
        frame_array = [cv2.cvtColor(np.array(img), cv2.COLOR_RGB2BGR) for img in cat_imgs]
        h, w, _ = frame_array[0].shape
        fourcc = cv2.VideoWriter_fourcc(*'mp4v')
        out = cv2.VideoWriter(video_path, fourcc, 1000/150, (w, h))  # 150ms一帧≈6.67fps
        for frame in frame_array:
            out.write(frame)
        out.release()
        print(f"视频已保存到: {video_path}")

        # 自动生成高质量视频
        temp_dir = osp.join(scene_dir, 'temp_frames')
        os.makedirs(temp_dir, exist_ok=True)
        for idx, img in enumerate(cat_imgs):
            img_even = pad_to_even(img)
            img_even.save(osp.join(temp_dir, f'frame_{idx:04d}.png'))
        hq_video_path = osp.join(scene_dir, 'video_resize_hq.mp4')
        ffmpeg_cmd = [
            'ffmpeg', '-y', '-framerate', f'{1000/150}',
            '-i', osp.join(temp_dir, 'frame_%04d.png'),
            '-c:v', 'libx264', '-crf', '10', '-pix_fmt', 'yuv420p',
            hq_video_path
        ]
        try:
            print('正在用ffmpeg生成高质量视频...')
            subprocess.run(ffmpeg_cmd, check=True)
            print(f'高质量视频已保存到: {hq_video_path}')
        except Exception as e:
            print(f'ffmpeg生成高质量视频失败: {e}')
        # 删除临时帧
        shutil.rmtree(temp_dir)
    
    